package com.lagou.edu.filter;
import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Bucket4j;
import io.github.bucket4j.Refill;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @Description: 自定义过滤器进行限流
 */
@Data
@Slf4j
@Component
public class RateLimitFilter implements GlobalFilter, Ordered {
    /**
     * 单机网关限流用一个ConcurrentHashMap来存储 bucket，
     * 如果是分布式集群限流的话，可以采用 Redis等分布式解决方案
     */
    private static final Map<String, Bucket> LOCAL_CACHE = new ConcurrentHashMap<>();
    /**
     * 桶的最大容量，即能装载 Token 的最大数量
     */
//    @Value("${limit.capacity}")
    private int capacity=10;
    /**
     * 每次 Token 补充量
     */
//    @Value("${limit.refillTokens}")
    private int refillTokens=1;
    /**
     *补充 Token 的时间间隔,秒数
     */
//    @Value("${limit.refillDuration}")
    private int refillDuration=10;

    private Bucket createNewBucket() {
        Refill refill = Refill.of(refillTokens, Duration.ofSeconds(refillDuration));
        Bandwidth limit = Bandwidth.classic(capacity, refill);
        return Bucket4j.builder().addLimit(limit).build();
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        if(!exchange.getRequest().getURI().getPath().startsWith("/user/register/")){
            // 合法请求，放行，执行后续的过滤器
            return chain.filter(exchange);
        }
        String ip = exchange.getRequest().getRemoteAddress().getAddress().getHostAddress();
        //若ip不存在则创建一个Bucket（令牌桶）
        Bucket bucket = LOCAL_CACHE.computeIfAbsent(ip, k -> createNewBucket());
        if (bucket.tryConsume(1)) {
            return chain.filter(exchange);
        } else {
            //当可用的令牌书为0是，进行限流返回429状态码
            exchange.getResponse().setStatusCode(HttpStatus.SEE_OTHER);
            String data = "您频繁进⾏注册，请求已被拒绝";
            DataBuffer wrap = exchange.getResponse().bufferFactory().wrap(data.getBytes());
            return exchange.getResponse().writeWith(Mono.just(wrap));
        }
    }

    @Override
    public int getOrder() {
        return -1000;
    }

    public static Map<String, Bucket> getLocalCache() {
        return LOCAL_CACHE;
    }


}